from global_var import *
from sklearn.model_selection import KFold

class FoodDataset(Dataset):
    def __init__(self, paths):
        super(FoodDataset).__init__()
        self.files = []
        for path in paths:
            self.files += [os.path.join(path,x) for x in os.listdir(path) if x.endswith(".jpg")]
        self.files = sorted(self.files)

    def __len__(self):
        return len(self.files)

    def __getitem__(self,idx):
        fname = self.files[idx]
        im = Image.open(fname)
        try:
            label = int(fname.split("/")[-1].split("_")[0])
        except:
            label = -1 # test has no label

        return im,label


class MixupDataset(Dataset):
    def __init__(self, dataset, alpha=0.8):
        self.dataset = dataset
        self.alpha = alpha

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        # 获取原始图像和标签
        img1, label1 = self.dataset[index]
        
        # 随机选择另一张图片和标签
        index2 = random.randint(0, len(self.dataset) - 1)
        img2, _ = self.dataset[index2]
        
        # Mixup图像（线性组合）
        img = self.alpha * img1 + (1 - self.alpha) * img2

        # Mixup标签（线性组合）
        return img, label1

class TransformDataset(Dataset):
    def __init__(self, dataset, tfm):
        super(TransformDataset).__init__()
        self.dataset = dataset
        self.transform = tfm
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        im, label = self.dataset[idx]
        return self.transform(im), label
    

# Construct train and valid datasets.
# The argument "loader" tells how torchvision reads the data.
full_dataset = FoodDataset([train_data_path, valid_data_path])
k_folds = 5
kfold = KFold(n_splits=k_folds, shuffle=True)
dataset_size = len(full_dataset)

# Construct test datasets.
# The argument "loader" tells how torchvision reads the data.
test_set = TransformDataset(FoodDataset([test_data_path]), test_tfm)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True)

test_tfm_set = TransformDataset(FoodDataset([test_data_path]), train_tfm)
test_tfm_loader = DataLoader(test_tfm_set, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True)